package main

import (
	"context"
	"github.com/olivere/elastic/v7"
)

var Client *elastic.Client

type EsConfig struct {
	Url      string `json:"url"`
	Sniff    bool   `json:"sniff"`
	Username string `json:"username"`
	Password string `json:"password"`
}

func (c EsConfig) Init() {
	var err error
	Client, err = elastic.NewClient(
		elastic.SetSniff(c.Sniff),
		elastic.SetURL(c.Url),
		elastic.SetBasicAuth(c.Username, c.Password))
	if err != nil {
		panic(err)
	}

	_, _, err = Client.Ping(c.Url).Do(context.Background())
	if err != nil {
		panic(err)
	}
}

// Insert 插入一条数据，如果插入id为空，则自动生成id
func Insert(ctx context.Context, index, id string, data []byte) error {
	s := Client.Index().Index(index)
	if id != "" {
		s = s.Id(id)
	}
	_, err := s.BodyJson(string(data)).Do(ctx)
	return err
}

// BatchInsert 批量插入,不指定id，自動生成id
func BatchInsert(ctx context.Context, index string, data [][]byte) error {
	bulkService := Client.Bulk().Index(index)

	for _, datum := range data {
		bulkService.Add(elastic.NewBulkIndexRequest().Doc(string(datum)))
	}
	_, err := bulkService.Do(ctx)
	return err
}

// BatchInsertById 批量插入,並指定id
func BatchInsertById(ctx context.Context, index string, dataMap map[string][]byte) error {
	bulkService := Client.Bulk().Index(index)

	for id, datum := range dataMap {
		bulkService.Add(elastic.NewBulkIndexRequest().Index(id).Doc(string(datum)))
	}
	_, err := bulkService.Do(ctx)
	return err
}

// Update
func Update(ctx context.Context, index, id string, dataMap map[string]interface{}) error {
	_, err := Client.Update().Index(index).Id(id).Doc(dataMap).Do(ctx)
	return err
}

func Delete(ctx context.Context, index, id string) error {
	_, err := Client.Delete().Index(index).Id(id).Do(ctx)
	return err
}

// DeleteBulk 删除文档批量（部分失败会返回失败）
// rets 删除的数量, error
func DeleteBulk(ctx context.Context, index string, ids []string) (int, error) {
	bulkService := Client.Bulk().Index(index).Refresh("true")
	for i := range ids {
		req := elastic.NewBulkDeleteRequest().Id(ids[i])
		bulkService.Add(req)
	}
	res, err := bulkService.Do(ctx)
	return len(res.Succeeded()), err
}

func GetById(index, id string) ([]byte, error) {
	result, err := Client.Get().Index(index).Id(id).Do(context.Background())
	if err != nil {
		return nil, err
	}
	if result != nil {
		return result.Source.MarshalJSON()
	}
	return nil, nil
}

func Search(index string, query elastic.Query) (map[string][]byte, int64, error) {
	result, err := Client.Search().Index(index).Query(query).Pretty(true).Do(context.Background())
	if err != nil {
		return nil, 0, err
	}

	if result.Hits.TotalHits.Value == 0 {
		return nil, 0, nil
	}

	bf := make(map[string][]byte)
	for _, hit := range result.Hits.Hits {
		bf[hit.Id] = hit.Source
	}
	return bf, result.Hits.TotalHits.Value, nil
}

func SearchBySort(index string, query elastic.Query, sorter elastic.Sorter) (map[string][]byte, int64, error) {
	result, err := Client.Search().Index(index).Query(query).SortBy(sorter).Do(context.Background())
	if err != nil {
		return nil, 0, err
	}
	if result.Hits.TotalHits.Value == 0 {
		return nil, 0, nil
	}

	bf := make(map[string][]byte)
	for _, hit := range result.Hits.Hits {
		bf[hit.Id] = hit.Source
	}
	return bf, result.Hits.TotalHits.Value, nil
}
